import cv2
import numpy as np
import dlib
import time
import os
from pydantic import BaseModel


class Response(BaseModel):
    image_path: str = ""
    center: tuple = (0, 0)
    face_width: int = 0
    face_height: int = 0
    image_wdith: int = 0
    image_height: int = 0
    point_8: tuple = (0, 0)  # 第8个点位置


class Face68:

    def __init__(self, model_path=None):
        # 加载人脸检测器和关键点检测器
        self.detector = dlib.get_frontal_face_detector()
        self.predictor = dlib.shape_predictor(
            "/models/.cache/shape_predictor_68_face_landmarks/shape_predictor_68_face_landmarks.dat"
            if not model_path else model_path)

    def get_keypoints_info(self, image_path, save_dir, target_face_width=None) -> Response:
        # 读取图像
        image = cv2.imread(image_path)
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 使用人脸检测器检测人脸
        faces = self.detector(gray)

        if len(faces) == 0:
            print("未检测到人脸")
            return None

        # 获取第一个人脸的关键点
        face = faces[0]
        landmarks = self.predictor(gray, face)
        landmarks = np.array([[p.x, p.y] for p in landmarks.parts()])

        # 计算人脸的中心坐标
        center_x = int(np.mean(landmarks[:, 0]))
        center_y = int(np.mean(landmarks[:, 1]))

        # 计算人脸的旋转角度
        angle = np.degrees(np.arctan2(landmarks[30, 1] - landmarks[27, 1],
                                      landmarks[30, 0] - landmarks[27, 0])) - 90

        # 构造旋转矩阵
        rotation_matrix = cv2.getRotationMatrix2D(
            (center_x, center_y), angle, 1)

        # 旋转图像
        aligned_image = cv2.warpAffine(
            image, rotation_matrix, (image.shape[1], image.shape[0]))

        face_width = face.width()
        face_height = face.height()
        center = (center_x, center_y)
        # 下颚中心
        point_8 = (int(landmarks[8, 0]), int(landmarks[8, 1]))
        height, width = aligned_image.shape[:2]

        if target_face_width:
            scale = face_width / target_face_width
            face_width = target_face_width
            face_height = int(face_height / scale)
            center = (int(center_x/scale), int(center_y/scale))
            point_8 = (int(point_8[0]/scale), int(point_8[1]/scale))
            height, width = int(height/scale), int(width/scale)
            aligned_image = cv2.resize(aligned_image, (width, height))
        filename = os.path.join(save_dir, "face-"+str(time.time())+".png")
        cv2.imwrite(filename, aligned_image)

        return Response(
            image_path=filename,
            center=center,
            face_width=face_width,
            face_height=face_height,
            image_wdith=width,
            image_height=height,
            point_8=point_8
        )


if __name__ == "__main__":
    model = Face68()
    res = model.get_keypoints_info(
        image_path="./origin.jpg", save_dir="./", target_face_width=400)
    print(res)
